package domain

import (
	"encoding/xml"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strconv"
)

//https://www.jianshu.com/p/dd8a113b02a3 go语言执行shell脚本

//https://www.cnblogs.com/caishunzhe/articles/12667889.html

const NvidiaSmiPath = "/usr/bin/nvidia-smi"

var testMode string

/**
GPU信息
*/
type Gpu struct {
	Temp        []GpuItem `json:"temp"`        //温度(C)
	Usage       []GpuItem `json:"usage"`       //显存使用(使用情况 MiB)
	UsedPercent []GpuItem `json:"usedPercent"` //使用占比(%)
}

//temp: [{ data:["72","72C"],name: 0 },{ data:["72","72C"],name: 1 }]
//usage: [{ data:["8493","8493Mib/15079Mib"],name: 0 },{ data:["9403","9403Mib/15079Mib"],name: 1 }]
//usedPercent: [{ data:["56","56%"],name: 0 },{ data:["43","43%"],name: 1 }]
type GpuItem struct {
	Data []string `json:"data"`
	Name int      `json:"name"`
}

/**
  获取gpu信息
*/
func (gpu *Gpu) GpuInfo() {

	//获得GPU metrics 信息
	metrics(gpu)

}

type NvidiaSmiLog struct {
	DriverVersion string `xml:"driver_version"`
	AttachedGPUs  string `xml:"attached_gpus"`
	GPUs          []struct {
		ProductName  string `xml:"product_name"`
		ProductBrand string `xml:"product_brand"`
		UUID         string `xml:"uuid"`
		FanSpeed     string `xml:"fan_speed"`
		PCI          struct {
			PCIBus string `xml:"pci_bus"`
		} `xml:"pci"`
		FbMemoryUsage struct {
			Total string `xml:"total"`
			Used  string `xml:"used"`
			Free  string `xml:"free"`
		} `xml:"fb_memory_usage"`
		Utilization struct {
			GPUUtil    string `xml:"gpu_util"`
			MemoryUtil string `xml:"memory_util"`
		} `xml:"utilization"`
		Temperature struct {
			GPUTemp              string `xml:"gpu_temp"`
			GPUTempMaxThreshold  string `xml:"gpu_temp_max_threshold"`
			GPUTempSlowThreshold string `xml:"gpu_temp_slow_threshold"`
		} `xml:"temperature"`
		PowerReadings struct {
			PowerDraw  string `xml:"power_draw"`
			PowerLimit string `xml:"power_limit"`
		} `xml:"power_readings"`
		Clocks struct {
			GraphicsClock string `xml:"graphics_clock"`
			SmClock       string `xml:"sm_clock"`
			MemClock      string `xml:"mem_clock"`
			VideoClock    string `xml:"video_clock"`
		} `xml:"clocks"`
		MaxClocks struct {
			GraphicsClock string `xml:"graphics_clock"`
			SmClock       string `xml:"sm_clock"`
			MemClock      string `xml:"mem_clock"`
			VideoClock    string `xml:"video_clock"`
		} `xml:"max_clocks"`
	} `xml:"gpu"`
}

func formatValue(key string, meta string, value string) string {
	result := key
	if meta != "" {
		result += "{" + meta + "}"
	}
	return result + " " + value + "\n"
}

func filterNumber(value string) string {
	r := regexp.MustCompile("[^0-9.]")
	return r.ReplaceAllString(value, "")
}

func metrics(gpu *Gpu) {

	var cmd *exec.Cmd
	if testMode == "1" {
		dir, err := os.Getwd()
		if err != nil {
			log.Fatal(err)
		}
		cmd = exec.Command("/bin/cat", dir+"/test.xml")
	} else {
		cmd = exec.Command(NvidiaSmiPath, "-q", "-x")
	}

	// Execute system command
	stdout, err := cmd.Output()
	if err != nil {

		//初始信息
		log.Println("解析GPU信息错误，错误信息为：", err.Error())
		//一块显卡的温度信息
		dataTemp := []string{"72", "72C"}
		itemTemp := GpuItem{Data: dataTemp, Name: 0}
		itemsTemp := []GpuItem{itemTemp}
		gpu.Temp = itemsTemp
		//一块显卡的显存使用信息
		dataUsage := []string{"8493", "8493Mib/15079Mib"}
		itemUsage := GpuItem{Data: dataUsage, Name: 0}
		itemsUsage := []GpuItem{itemUsage}
		gpu.Usage = itemsUsage
		//一块显卡的使用占比信息
		dataUsedPercent := []string{"56", "56%"}
		itemUsedPercent := GpuItem{Data: dataUsedPercent, Name: 0}
		itemsUsedPercent := []GpuItem{itemUsedPercent}
		gpu.UsedPercent = itemsUsedPercent

	} else {

		//一块显卡的温度信息

		// Parse XML
		var xmlData NvidiaSmiLog
		_ = xml.Unmarshal(stdout, &xmlData)

		//创建GPU温度切片
		itemsTemp := make([]GpuItem, len(xmlData.GPUs))

		//创建显卡的显存使用信息切片
		itemsUsage := make([]GpuItem, len(xmlData.GPUs))

		//创建显卡的使用占比信息切片
		itemsUsedPercent := make([]GpuItem, len(xmlData.GPUs))

		log.Println("nvidiasmi_driver_version:%s", xmlData.DriverVersion)

		log.Println("nvidiasmi_attached_gpus:%s", filterNumber(xmlData.AttachedGPUs))

		for index, info := range xmlData.GPUs {

			//温度信息开始
			gpuTemp := filterNumber(info.Temperature.GPUTemp)

			dataTemp := []string{gpuTemp, gpuTemp + "C"}

			itemsTemp[index] = GpuItem{Data: dataTemp, Name: index}
			//温度信息结束

			//显存使用情况开始
			used := filterNumber(info.FbMemoryUsage.Used)
			total := filterNumber(info.FbMemoryUsage.Total)

			dataUsage := []string{used, used + "Mib/" + total + "Mib"}

			itemsUsage[index] = GpuItem{Data: dataUsage, Name: index}

			//显存使用情况结束
			usedInt, _ := strconv.ParseInt(used, 10, 64)
			totalInt, _ := strconv.ParseInt(used, 10, 64)

			UsedPercent := float32(0.0)
			if usedInt == 0 || totalInt == 0 {
				UsedPercent = 0.00
			} else {
				UsedPercent = float32(usedInt / totalInt)
			}

			//显卡的使用占比信息开始
			dataUsedPercent := []string{FloatToString(UsedPercent), (FloatToString(UsedPercent)) + "%"}
			itemsUsedPercent[index] = GpuItem{Data: dataUsedPercent, Name: index}
			//显卡的使用占比信息结束

		}

		gpu.Temp = itemsTemp
		gpu.Usage = itemsUsage
		gpu.UsedPercent = itemsUsedPercent

	}

}

//float32 转 String工具类，保留6位小数
func FloatToString(input_num float32) string {
	// to convert a float number to a string
	return strconv.FormatFloat(float64(input_num), 'f', 2, 64)
}
